#ifndef MODELTORCH
#define MODELTORCH

#include <memory>
#include <exception>
#include <vector>
#include <string>
#include <vector>
#include <torch/script.h>
#include <torch/cuda.h> 
#include <opencv2/opencv.hpp>
#include <torch/torch.h>

namespace Bublik
{

 class ModelTorch
 {
 private:
 torch::NoGradGuard no_grad;
 torch::jit::script::Module _model;
 cv::Mat _inputImage;
 c10::DeviceType _device;
 int _imgSize_x;
 int _imgSize_y;
 cv::Mat PrepImg(cv::Mat& input_img, int inputX, int inputY);

 public:
 ModelTorch(std::string& FileName);
 
 cv::Mat getDetection(std::string& ImageFileName, float conf);
 cv::Mat getSegmentation(std::string& ImageFileName, float conf);

 protected:

 };

} //namespase Bublik {

#endif